import numpy as np

def calculate_iou(box1, box2):
    """
    计算两个边界框的IoU (Intersection over Union)。
    box: [x1, y1, x2, y2]
    """
    # 提取坐标
    x1_b1, y1_b1, x2_b1, y2_b1 = box1
    x1_b2, y1_b2, x2_b2, y2_b2 = box2

    # 计算交集区域的坐标
    x1_inter = max(x1_b1, x1_b2)
    y1_inter = max(y1_b1, y1_b2)
    x2_inter = min(x2_b1, x2_b2)
    y2_inter = min(y2_b1, y2_b2)

    # 计算交集区域的面积
    inter_width = max(0, x2_inter - x1_inter)
    inter_height = max(0, y2_inter - y1_inter)
    inter_area = inter_width * inter_height

    # 计算两个边界框的面积
    area_b1 = (x2_b1 - x1_b1) * (y2_b1 - y1_b1)
    area_b2 = (x2_b2 - x1_b2) * (y2_b2 - y1_b2)

    # 计算并集面积
    union_area = area_b1 + area_b2 - inter_area

    if union_area == 0:
        return 0.0
    
    iou = inter_area / union_area
    return iou

def calculate_diou(box1, box2):
    """
    计算两个边界框的DIoU (Distance-IoU)。
    box: [x1, y1, x2, y2]
    """
    iou = calculate_iou(box1, box2)

    # 提取坐标
    x1_b1, y1_b1, x2_b1, y2_b1 = box1
    x1_b2, y1_b2, x2_b2, y2_b2 = box2

    # 计算中心点坐标
    center_x1 = (x1_b1 + x2_b1) / 2
    center_y1 = (y1_b1 + y2_b1) / 2
    center_x2 = (x1_b2 + x2_b2) / 2
    center_y2 = (y1_b2 + y2_b2) / 2

    # 计算中心点欧氏距离的平方
    distance_squared = (center_x1 - center_x2)**2 + (center_y1 - center_y2)**2

    # 计算最小外接矩形C的坐标
    x1_c = min(x1_b1, x1_b2)
    y1_c = min(y1_b1, y1_b2)
    x2_c = max(x2_b1, x2_b2)
    y2_c = max(y2_b1, y2_b2)

    # 计算最小外接矩形C的对角线长度的平方
    c_diag_squared = (x2_c - x1_c)**2 + (y2_c - y1_c)**2

    # 避免除以零，特别是当c_diag_squared为0时 (例如，两个框都是点)
    if c_diag_squared == 0:
        # 如果外接框对角线为0，通常表示框退化为点，此时DIoU设为-1或0，表示无法有效计算
        # 这里设置为和IoU一样（如果IoU也不为0，但这种情况一般不会出现）
        # 或者直接返回IoU，因为惩罚项为0
        return iou if iou > 0 else -1.0 # 如果IoU为0，表示完全不重叠，此时可以视为无法有效计算DIoU的惩罚项
                                     # 这里为了演示，我们让它表现出负值或特殊值
                                     # 实际应用中需要根据具体情况处理
    
    # 计算距离惩罚项
    distance_penalty = distance_squared / c_diag_squared

    diou = iou - distance_penalty
    return diou

if __name__ == "__main__":
    # 定义一些示例边界框
    # box格式: [x1, y1, x2, y2]

    # 场景1：两个框完全重合 (期待IoU=1, DIoU=1)
    box_a = [0, 0, 10, 10]
    box_b = [0, 0, 10, 10]
    iou_ab = calculate_iou(box_a, box_b)
    diou_ab = calculate_diou(box_a, box_b)
    print(f"场景1：框A {box_a}, 框B {box_b} (完全重合)")
    print(f"IoU: {iou_ab:.4f}, DIoU: {diou_ab:.4f}\n") # DIoU和IoU都为1

    # 场景2：两个框部分重叠，但中心点有偏差 (期待IoU>0, DIoU < IoU)
    box_c = [0, 0, 10, 10]
    box_d = [2, 2, 12, 12] # 中心点有偏差
    iou_cd = calculate_iou(box_c, box_d)
    diou_cd = calculate_diou(box_c, box_d)
    print(f"场景2：框C {box_c}, 框D {box_d} (部分重叠，中心点有偏差)")
    print(f"IoU: {iou_cd:.4f}, DIoU: {diou_cd:.4f}\n") # DIoU会更低，因为它惩罚了中心点距离

    # 场景3：两个框没有交集，但中心点在一条直线上 (期待IoU=0, DIoU为负值，能区分距离)
    box_e = [0, 0, 5, 5]
    box_f = [6, 0, 11, 5] # 相距1个单位
    iou_ef = calculate_iou(box_e, box_f)
    diou_ef = calculate_diou(box_e, box_f)
    print(f"场景3：框E {box_e}, 框F {box_f} (无交集，中心点对齐)")
    print(f"IoU: {iou_ef:.4f}, DIoU: {diou_ef:.4f}\n") # IoU为0，DIoU为负值

    # 场景4：两个框没有交集，且离得很远 (期待IoU=0, DIoU负值更小)
    box_g = [0, 0, 5, 5]
    box_h = [100, 100, 105, 105] # 相距很远
    iou_gh = calculate_iou(box_g, box_h)
    diou_gh = calculate_diou(box_g, box_h)
    print(f"场景4：框G {box_g}, 框H {box_h} (无交集，离得很远)")
    print(f"IoU: {iou_gh:.4f}, DIoU: {diou_gh:.4f}\n") # IoU为0，DIoU负值会更小

    # 场景5：一个框完全包含另一个框 (期待IoU>0, DIoU接近IoU但会惩罚中心点偏差)
    box_i = [0, 0, 20, 20] # 大框
    box_j = [5, 5, 15, 15] # 小框，被包含
    iou_ij = calculate_iou(box_i, box_j)
    diou_ij = calculate_diou(box_i, box_j)
    print(f"场景5：框I {box_i} (大框), 框J {box_j} (被包含)")
    print(f"IoU: {iou_ij:.4f}, DIoU: {diou_ij:.4f}\n") # DIoU会小于IoU，因为中心点有偏差

    # 场景6：同心但大小不同的框 (期待IoU>0, DIoU=IoU，因为中心点距离为0)
    box_k = [0, 0, 20, 20]
    box_l = [5, 5, 15, 15] # 中心点与K相同
    iou_kl = calculate_iou(box_k, box_l)
    diou_kl = calculate_diou(box_k, box_l)
    print(f"场景6：框K {box_k}, 框L {box_l} (同心但大小不同)")
    print(f"IoU: {iou_kl:.4f}, DIoU: {diou_kl:.4f}\n") # DIoU=IoU，因为中心点距离惩罚项为0